#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
将Redis数据导入到Hive
"""

import redis
import json
import subprocess
import tempfile
import os
from datetime import datetime

def connect_redis():
    """连接Redis"""
    r = redis.Redis(host='192.168.88.130', port=6379, password='123456', decode_responses=True)
    r.ping()
    return r

def export_to_csv():
    """导出数据到CSV文件"""
    r = connect_redis()
    keys = r.keys("cnhnb_price:*")
    keys.sort()
    
    # 创建临时CSV文件
    csv_file = f"cnhnb_data_{datetime.now().strftime('%Y%m%d_%H%M%S')}.csv"
    
    with open(csv_file, 'w', encoding='utf-8') as f:
        # 写入CSV头
        f.write("name,price_value,price_original,unit,location,date,timestamp,crawl_timestamp,source_url,spider_name,dt\n")
        
        for key in keys:
            try:
                raw_data = r.get(key)
                if raw_data:
                    data = json.loads(raw_data)
                    
                    # 提取日期作为分区
                    crawl_time = data.get('crawl_timestamp', '')
                    if crawl_time:
                        dt = crawl_time.split(' ')[0] if ' ' in crawl_time else datetime.now().strftime('%Y-%m-%d')
                    else:
                        dt = datetime.now().strftime('%Y-%m-%d')
                    
                    # 写入CSV行
                    f.write(f"\"{data.get('name', '')}\",")
                    f.write(f"{data.get('price_value', '')},")
                    # 处理原始价格文本中的引号
                    price_original = data.get('price_original', '').replace('"', '""')
                    f.write(f"\"{price_original}\",")
                    f.write(f"\"{data.get('unit', '')}\",")
                    f.write(f"\"{data.get('location', '')}\",")
                    f.write(f"\"{data.get('date', '')}\",")
                    f.write(f"{data.get('timestamp', '')},")
                    f.write(f"\"{data.get('crawl_timestamp', '')}\",")
                    f.write(f"\"{data.get('source_url', '')}\",")
                    f.write(f"\"{data.get('spider_name', '')}\",")
                    f.write(f"\"{dt}\"\n")
                    
            except Exception as e:
                print(f"处理键 {key} 时出错: {e}")
    
    print(f"✓ 数据已导出到CSV文件: {csv_file}")
    return csv_file

def create_hive_table():
    """创建Hive表"""
    hive_sql = """
    -- 创建数据库
    CREATE DATABASE IF NOT EXISTS agricultural_data;
    
    -- 使用数据库
    USE agricultural_data;
    
    -- 创建价格数据表
    CREATE TABLE IF NOT EXISTS cnhnb_prices (
        name STRING COMMENT '产品名称',
        price_value DOUBLE COMMENT '价格数值',
        price_original STRING COMMENT '原始价格字符串',
        unit STRING COMMENT '价格单位',
        location STRING COMMENT '产地/地区',
        `date` STRING COMMENT '价格日期',
        `timestamp` DOUBLE COMMENT 'Unix时间戳',
        crawl_timestamp STRING COMMENT '爬取时间',
        source_url STRING COMMENT '数据来源URL',
        spider_name STRING COMMENT '爬虫名称'
    )
    COMMENT '惠农网农产品价格数据表'
    PARTITIONED BY (dt STRING COMMENT '分区日期 YYYY-MM-DD')
    ROW FORMAT DELIMITED
    FIELDS TERMINATED BY ','
    STORED AS TEXTFILE;
    """
    
    # 保存SQL到文件
    sql_file = "create_hive_table.sql"
    with open(sql_file, 'w', encoding='utf-8') as f:
        f.write(hive_sql)
    
    print(f"✓ Hive建表SQL已保存到: {sql_file}")
    return sql_file

def load_to_hive(csv_file, partition_date=None):
    """加载数据到Hive"""
    if not partition_date:
        partition_date = datetime.now().strftime('%Y-%m-%d')
    
    # 创建Hive加载SQL
    load_sql = f"""
    USE agricultural_data;
    
    -- 加载数据到Hive
    LOAD DATA LOCAL INPATH '{csv_file}' 
    OVERWRITE INTO TABLE cnhnb_prices 
    PARTITION (dt='{partition_date}');
    
    -- 验证数据
    SELECT COUNT(*) as total_records FROM cnhnb_prices WHERE dt='{partition_date}';
    SELECT * FROM cnhnb_prices WHERE dt='{partition_date}' LIMIT 5;
    """
    
    # 保存加载SQL
    load_file = "load_to_hive.sql"
    with open(load_file, 'w', encoding='utf-8') as f:
        f.write(load_sql)
    
    print(f"✓ Hive加载SQL已保存到: {load_file}")
    print(f"✓ 分区日期: {partition_date}")
    
    return load_file

def main():
    print("=== 开始将Redis数据导入Hive ===")
    
    try:
        # 1. 导出数据到CSV
        csv_file = export_to_csv()
        
        # 2. 创建Hive表
        create_hive_table()
        
        # 3. 生成加载SQL
        load_to_hive(csv_file)
        
        print("\n=== 导入准备完成 ===")
        print("请按以下步骤执行：")
        print("1. 执行建表SQL:")
        print("   hive -f create_hive_table.sql")
        print("2. 执行数据加载SQL:")
        print("   hive -f load_to_hive.sql")
        print("3. 验证数据:")
        print("   hive -e 'USE agricultural_data; SELECT COUNT(*) FROM cnhnb_prices;'")
        
    except Exception as e:
        print(f"✗ 导入过程出错: {e}")

if __name__ == '__main__':
    main() 